#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <iostream>
#include <fstream>
#include <unistd.h>
#include <algorithm> 
#include <vector>
#include <inttypes.h>

#include "libpsacyber.h"

#define TEST_OBJECT_SIZE	7000

void cout_hex(uint8_t *addr, size_t len)
{
    if(addr == nullptr)
        return;
    for(size_t i = 0; i < len; i++) {
        printf("%02x", addr[i]);
    }
    std::cout << std::endl;
}

static void print_hex(const char *info, uint8_t*data, size_t len){
    printf("%s[len:%zu | 0x%zx]\n",info,len,len);
    for(size_t i=0;i<len;i++){
        printf("%02x ",data[i]);
    }
    printf("\n");
}

/* Func: test_ftcm_hsm_ap_generate_random
 *
 * TEST:
 *   - ftcm_hsm_ap_generate_random: used to generating a random number
 */
void test_ftcm_hsm_ap_generate_random()
{
    psa_status_t ret = PSA_OK;

    uint8_t buf[4] = {0};
    size_t output_size = 4;

    printf("\n-----\nTEST ftcm_hsm_ap_generate_random():\n");
    printf("INPUT: buf=%s output_size=%zu\n", buf, output_size);

    ret = ftcm_hsm_ap_generate_random(buf, output_size);

    printf("OUTPUT: ret=%d buf=0x", ret);
    print_hex("Random data", buf, output_size);
}

/* Func: test_ftcm_hsm_ap_hash_compute
 *
 * TEST:
 *   - ftcm_hsm_ap_hash_compute: used to compute hash value
 */
void test_ftcm_hsm_ap_hash_compute()
{
    psa_status_t ret = PSA_OK;

    uint8_t data[] = "abc123";
    size_t dataLen = 6;
    uint8_t hash[100] = {0};
    size_t hash_size = 100;
    size_t hash_length = 0;

    // ftcm_hsm_ap_hash_compute with SHA256
    printf("\n-----\nTEST ftcm_hsm_ap_hash_compute():\n");
    printf("INPUT: data=%s dataLen=%zu alg=PSA_ALG_SHA_256 hash=%s hash_size=%zu\n", data, dataLen, hash, hash_size);

    ret = ftcm_hsm_ap_hash_compute(PSA_ALG_SHA_256, data, dataLen, hash, hash_size, &hash_length);

    printf("OUTPUT: ret=%d hash=0x", ret);
    print_hex("Hash", hash, hash_length);
    printf(" hash_length=%zu\n", hash_length);
}

/* Func: test_ftcm_hsm_ap_hash_operations
 *
 * TEST:
 *   - ftcm_hsm_ap_hash_setup: used to setup hash operation
 *   - ftcm_hsm_ap_hash_update: used to update hash operation
 *   - ftcm_hsm_ap_hash_finish: used to finish hash operation
 */
void test_ftcm_hsm_ap_hash_operations()
{
    psa_status_t ret = PSA_OK;


    uint8_t data[] = "abc123";
    size_t dataLen = 6;
    uint8_t hash[100] = {0};
    size_t hash_size = 100;
    size_t hash_length = 0;

    // ftcm_hsm_ap_hash_setup
    printf("\n-----\nTEST ftcm_hsm_ap_hash_setup/update/finish():\n");
    printf("INPUT: alg=PSA_ALG_SHA_256\n");

    ret = ftcm_hsm_ap_hash_setup(NULL, PSA_ALG_SHA_256);
    printf("ftcm_hsm_ap_hash_setup OUTPUT: ret=%d\n", ret);

    // ftcm_hsm_ap_hash_update
    ret = ftcm_hsm_ap_hash_update(NULL, data, dataLen);
    printf("ftcm_hsm_ap_hash_update OUTPUT: ret=%d\n", ret);

    // ftcm_hsm_ap_hash_finish
    ret = ftcm_hsm_ap_hash_finish(NULL, hash, hash_size, &hash_length);
    printf("ftcm_hsm_ap_hash_finish OUTPUT: ret=%d hash=0x", ret);
    print_hex("Hash operation result", hash, hash_length);
    printf(" hash_length=%zu\n", hash_length);
}

/* Func: test_ftcm_hsm_ap_cipher_encrypt_decrypt
 *
 * TEST:
 *   - ftcm_hsm_ap_cipher_encrypt: used to encrypt data
 *   - ftcm_hsm_ap_cipher_decrypt: used to decrypt data
 */
void test_ftcm_hsm_ap_cipher_encrypt_decrypt()
{
    psa_status_t ret = PSA_OK;
    mbedtls_svc_key_id_t key = 0x12340002; // 示例密钥ID

    uint8_t data[16] = "abcdefgh1234567";
    size_t dataLen = 16;
    uint8_t encrypted[32] = {0}; // 加密结果
    uint8_t decrypted[32] = {0}; // 解密结果
    size_t output_length = 0;
    size_t output_size = 32;

    // ftcm_hsm_ap_cipher_encrypt
    printf("\n-----\nTEST ftcm_hsm_ap_cipher_encrypt():\n");
    printf("INPUT: key=%d alg=PSA_ALG_CBC_NO_PADDING data=%s dataLen=%zu\n", key, data, dataLen);

    ret = ftcm_hsm_ap_cipher_encrypt(key, PSA_ALG_CBC_NO_PADDING, data, dataLen, 
                                     encrypted, output_size, &output_length);

    printf("OUTPUT: ret=%d encrypted=0x", ret);
    print_hex("Encrypted data", encrypted, output_length);
    printf(" output_length=%zu\n", output_length);

    // ftcm_hsm_ap_cipher_decrypt
    printf("\n-----\nTEST ftcm_hsm_ap_cipher_decrypt():\n");
    printf("INPUT: key=%d alg=PSA_ALG_CBC_NO_PADDING encrypted=... encryptedLen=%zu\n", key, output_length);

    ret = ftcm_hsm_ap_cipher_decrypt(key, PSA_ALG_CBC_NO_PADDING, encrypted, output_length, 
                                     decrypted, output_size, &output_length);

    printf("OUTPUT: ret=%d decrypted=%s output_length=%zu\n", ret, decrypted, output_length);
}

/* Func: test_ftcm_hsm_ap_mac_compute
 *
 * TEST:
 *   - ftcm_hsm_ap_mac_compute: used to compute MAC value
 */
void test_ftcm_hsm_ap_mac_compute()
{
    psa_status_t ret = PSA_OK;
    mbedtls_svc_key_id_t key = 0x12340002; // 示例密钥ID

    uint8_t data[] = "abc123";
    size_t data_len = 6;
    uint8_t mac[64] = {0};
    size_t mac_size = 64;
    size_t mac_length = 0;

    // ftcm_hsm_ap_mac_compute
    printf("\n-----\nTEST ftcm_hsm_ap_mac_compute():\n");
    printf("INPUT: key=%d alg=PSA_ALG_CMAC data=%s data_len=%zu mac_size=%zu\n", key, data, data_len, mac_size);

    ret = ftcm_hsm_ap_mac_compute(key, PSA_ALG_CMAC, data, data_len, mac, mac_size, &mac_length);

    printf("OUTPUT: ret=%d mac=0x", ret);
    print_hex("MAC", mac, mac_length);
    printf(" mac_length=%zu\n", mac_length);
}

/* Func: test_ftcm_hsm_ap_mac_operations
 *
 * TEST:
 *   - ftcm_hsm_ap_mac_setup: used to setup MAC operation
 *   - ftcm_hsm_ap_mac_update: used to update MAC operation
 *   - ftcm_hsm_ap_mac_finish: used to finish MAC operation
 */
void test_ftcm_hsm_ap_mac_operations()
{
    psa_status_t ret = PSA_OK;
    mac_compute_ctx_t ctx;
    mbedtls_svc_key_id_t key = 0x12340002; // 示例密钥ID

    uint8_t data[] = "abc123";
    size_t data_len = 6;
    uint8_t mac[64] = {0};
    size_t mac_size = 64;
    size_t mac_length = 0;

    // ftcm_hsm_ap_mac_setup
    printf("\n-----\nTEST ftcm_hsm_ap_mac_setup/update/finish():\n");
    printf("INPUT: key=%d alg=PSA_ALG_CMAC\n", key);

    ret = ftcm_hsm_ap_mac_setup(&ctx, key, PSA_ALG_CMAC);
    printf("ftcm_hsm_ap_mac_setup OUTPUT: ret=%d\n", ret);

    // ftcm_hsm_ap_mac_update
    ret = ftcm_hsm_ap_mac_update(&ctx, data, data_len);
    printf("ftcm_hsm_ap_mac_update OUTPUT: ret=%d\n", ret);

    // ftcm_hsm_ap_mac_finish
    ret = ftcm_hsm_ap_mac_finish(&ctx, mac, mac_size, &mac_length);
    printf("ftcm_hsm_ap_mac_finish OUTPUT: ret=%d mac=0x", ret);
    print_hex("MAC operation result", mac, mac_length);
    printf(" mac_length=%zu\n", mac_length);
}

/* Func: test_ftcm_hsm_ap_rsa_operations
 *
 * TEST:
 *   - ftcm_hsm_ap_rsa_generate_key: used to generate RSA key pair
 *   - ftcm_hsm_ap_rsa_encrypt: used to encrypt with RSA
 *   - ftcm_hsm_ap_rsa_decrypt: used to decrypt with RSA
 *   - ftcm_hsm_ap_rsa_sign: used to sign with RSA
 *   - ftcm_hsm_ap_rsa_verify: used to verify RSA signature
 */
void test_ftcm_hsm_ap_rsa_operations()
{
    psa_status_t ret = PSA_OK;
    mbedtls_svc_key_id_t key = 0;
    
    // For encryption/decryption
    uint8_t data[] = "abc123";
    size_t data_len = 6;
    uint8_t encrypted[512] = {0};
    size_t encrypted_size = 512;
    size_t encrypted_length = 0;
    uint8_t decrypted[512] = {0};
    size_t decrypted_size = 512;
    size_t decrypted_length = 0;
    
    // For sign/verify
    uint8_t signature[512] = {0};
    size_t signature_size = 512;
    size_t signature_length = 0;

    // Note: 这里只是示例测试调用，实际使用需要先生成密钥
    printf("\n-----\nTEST ftcm_hsm_ap_rsa operations():\n");
    
    // ftcm_hsm_ap_rsa_encrypt
    printf("INPUT: key=%d alg=PSA_ALG_RSA_PKCS1V15_SIGN_RAW\n", key);
    ret = ftcm_hsm_ap_rsa_encrypt(key, PSA_ALG_RSA_PKCS1V15_SIGN_RAW, data, data_len, 
                                  NULL, 0, encrypted, encrypted_size, &encrypted_length);
    printf("ftcm_hsm_ap_rsa_encrypt OUTPUT: ret=%d encrypted_length=%zu\n", ret, encrypted_length);

    // ftcm_hsm_ap_rsa_decrypt
    ret = ftcm_hsm_ap_rsa_decrypt(key, PSA_ALG_RSA_PKCS1V15_SIGN_RAW, encrypted, encrypted_length,
                                  NULL, 0, decrypted, decrypted_size, &decrypted_length);
    printf("ftcm_hsm_ap_rsa_decrypt OUTPUT: ret=%d decrypted_length=%zu\n", ret, decrypted_length);

    // ftcm_hsm_ap_rsa_sign
    ret = ftcm_hsm_ap_rsa_sign(key, PSA_ALG_RSA_PKCS1V15_SIGN_RAW, data, data_len,
                               signature, signature_size, &signature_length);
    printf("ftcm_hsm_ap_rsa_sign OUTPUT: ret=%d signature_length=%zu\n", ret, signature_length);

    // ftcm_hsm_ap_rsa_verify
    ret = ftcm_hsm_ap_rsa_verify(key, PSA_ALG_RSA_PKCS1V15_SIGN_RAW, data, data_len,
                                 signature, signature_length);
    printf("ftcm_hsm_ap_rsa_verify OUTPUT: ret=%d\n", ret);
}

/* Func: test_ftcm_hsm_ap_aead_operations
 *
 * TEST:
 *   - ftcm_hsm_ap_aead_encrypt: used to AEAD encrypt
 *   - ftcm_hsm_ap_aead_decrypt: used to AEAD decrypt
 */
void test_ftcm_hsm_ap_aead_operations()
{
    psa_status_t ret = PSA_OK;
    mbedtls_svc_key_id_t key = 0x12340002; // 示例密钥ID
    
    // Test data
    uint8_t plaintext[] = "secret message";
    size_t plaintext_length = strlen((char*)plaintext);
    uint8_t nonce[12] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b};
    size_t nonce_length = 12;
    size_t tag_length = 16;
    uint8_t additional_data[] = "header";
    size_t additional_data_length = strlen((char*)additional_data);
    
    uint8_t ciphertext[1024] = {0};
    size_t ciphertext_size = 1024;
    size_t ciphertext_length = 0;
    
    uint8_t decrypted[1024] = {0};
    size_t plaintext_size = 1024;
    size_t plaintext_length_out = 0;

    // ftcm_hsm_ap_aead_encrypt
    printf("\n-----\nTEST ftcm_hsm_ap_aead_encrypt():\n");
    printf("INPUT: key=%d alg=PSA_ALG_CMAC plaintext=%s plaintext_length=%zu\n", key, plaintext, plaintext_length);

    ret = ftcm_hsm_ap_aead_encrypt(key, PSA_ALG_CMAC, nonce, nonce_length, tag_length,
                                   additional_data, additional_data_length,
                                   plaintext, plaintext_length,
                                   ciphertext, ciphertext_size, &ciphertext_length);
    printf("ftcm_hsm_ap_aead_encrypt OUTPUT: ret=%d ciphertext_length=%zu\n", ret, ciphertext_length);

    // ftcm_hsm_ap_aead_decrypt
    printf("\n-----\nTEST ftcm_hsm_ap_aead_decrypt():\n");
    ret = ftcm_hsm_ap_aead_decrypt(key, PSA_ALG_CMAC, nonce, nonce_length, tag_length,
                                   additional_data, additional_data_length,
                                   ciphertext, ciphertext_length,
                                   decrypted, plaintext_size, &plaintext_length_out);
    printf("ftcm_hsm_ap_aead_decrypt OUTPUT: ret=%d plaintext_length_out=%zu\n", ret, plaintext_length_out);
    if (ret == PSA_SUCCESS && plaintext_length_out <= sizeof(decrypted)) {
        decrypted[plaintext_length_out] = '\0';
        printf("Decrypted text: %s\n", decrypted);
    }
}


int main(int argc, char **argv)
{
    std::cout << "|----------------------|" << std::endl;
    std::cout << "| libpsacyber APIs test |" << std::endl;
    std::cout << "|-----------------------|" << std::endl;

    test_ftcm_hsm_ap_generate_random();

    test_ftcm_hsm_ap_hash_compute();

    test_ftcm_hsm_ap_hash_operations();

    test_ftcm_hsm_ap_cipher_encrypt_decrypt();

    test_ftcm_hsm_ap_mac_compute();

    test_ftcm_hsm_ap_mac_operations();

    test_ftcm_hsm_ap_rsa_operations();

    test_ftcm_hsm_ap_aead_operations();

    return 0;
}